{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 02: Navier-Stokes problem\n",
    "\n",
    "In this tutorial we compare the formulation of a Navier-Stokes problem using standard assembly with mixed function spaces and block assembly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import (derivative, div, dx, FiniteElement, grad, inner, MixedElement, split, TestFunction,\n",
    "                 TrialFunction, VectorElement)\n",
    "from dolfinx import DirichletBC, Function, FunctionSpace\n",
    "from dolfinx.fem import locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (apply_lifting, assemble_matrix, assemble_matrix_block, assemble_scalar,\n",
    "                               assemble_vector, assemble_vector_block, BlockVecSubVectorWrapper,\n",
    "                               create_matrix, create_matrix_block, create_vector, create_vector_block,\n",
    "                               set_bc)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constitutive parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nu = 0.01\n",
    "\n",
    "\n",
    "def u_in_eval(x):\n",
    "    values = np.zeros((2, x.shape[1]))\n",
    "    values[0, :] = 1.0\n",
    "    return values\n",
    "\n",
    "\n",
    "def u_wall_eval(x):\n",
    "    return np.zeros((2, x.shape[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/backward_facing_step.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh()\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "boundaries_1 = boundaries.indices[boundaries.values == 1]\n",
    "boundaries_2 = boundaries.indices[boundaries.values == 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V_element = VectorElement(\"Lagrange\", mesh.ufl_cell(), 2)\n",
    "Q_element = FiniteElement(\"Lagrange\", mesh.ufl_cell(), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard FEniCSx formulation using a mixed function space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_monolithic():\n",
    "    # Function spaces\n",
    "    W_element = MixedElement(V_element, Q_element)\n",
    "    W = FunctionSpace(mesh, W_element)\n",
    "\n",
    "    # Test and trial functions: monolithic\n",
    "    vq = TestFunction(W)\n",
    "    (v, q) = split(vq)\n",
    "    dup = TrialFunction(W)\n",
    "    up = Function(W)\n",
    "    (u, p) = split(up)\n",
    "\n",
    "    # Variational forms\n",
    "    F = (nu * inner(grad(u), grad(v)) * dx\n",
    "         + inner(grad(u) * u, v) * dx\n",
    "         - div(v) * p * dx\n",
    "         + div(u) * q * dx)\n",
    "    J = derivative(F, up, dup)\n",
    "\n",
    "    # Boundary conditions\n",
    "    u_in = Function(W.sub(0).collapse())\n",
    "    u_in.interpolate(u_in_eval)\n",
    "    u_wall = Function(W.sub(0).collapse())\n",
    "    u_wall.interpolate(u_wall_eval)\n",
    "    bdofs_V_1 = locate_dofs_topological((W.sub(0), W.sub(0).collapse()), mesh.topology.dim - 1, boundaries_1)\n",
    "    bdofs_V_2 = locate_dofs_topological((W.sub(0), W.sub(0).collapse()), mesh.topology.dim - 1, boundaries_2)\n",
    "    inlet_bc = DirichletBC(u_in, bdofs_V_1, W.sub(0))\n",
    "    wall_bc = DirichletBC(u_wall, bdofs_V_2, W.sub(0))\n",
    "    bc = [inlet_bc, wall_bc]\n",
    "\n",
    "    # Class for interfacing with SNES\n",
    "    class NavierStokesProblem(object):\n",
    "        def __init__(self, F, J, solution, bcs, P=None):\n",
    "            self._F = F\n",
    "            self._J = J\n",
    "            self._obj_vec = create_vector(F)\n",
    "            self._solution = solution\n",
    "            self._bcs = bcs\n",
    "            self._P = P\n",
    "\n",
    "        def update_solution(self, x):\n",
    "            x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "            with x.localForm() as _x, self._solution.vector.localForm() as _solution:\n",
    "                _solution[:] = _x\n",
    "\n",
    "        def obj(self, snes, x):\n",
    "            self.F(snes, x, self._obj_vec)\n",
    "            return self._obj_vec.norm()\n",
    "\n",
    "        def F(self, snes, x, F_vec):\n",
    "            self.update_solution(x)\n",
    "            with F_vec.localForm() as F_vec_local:\n",
    "                F_vec_local.set(0.0)\n",
    "            assemble_vector(F_vec, self._F)\n",
    "            apply_lifting(F_vec, [self._J], [self._bcs], x0=[x], scale=-1.0)\n",
    "            F_vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)\n",
    "            set_bc(F_vec, self._bcs, x, -1.0)\n",
    "\n",
    "        def J(self, snes, x, J_mat, P_mat):\n",
    "            J_mat.zeroEntries()\n",
    "            assemble_matrix(J_mat, self._J, self._bcs, diagonal=1.0)\n",
    "            J_mat.assemble()\n",
    "            if self._P is not None:\n",
    "                P_mat.zeroEntries()\n",
    "                assemble_matrix(P_mat, self._P, self._bcs, diagonal=1.0)\n",
    "                P_mat.assemble()\n",
    "\n",
    "    # Create problem\n",
    "    problem = NavierStokesProblem(F, J, up, bc)\n",
    "    F_vec = create_vector(F)\n",
    "    J_mat = create_matrix(J)\n",
    "\n",
    "    # Solve\n",
    "    snes = PETSc.SNES().create(mesh.mpi_comm())\n",
    "    snes.setTolerances(max_it=20)\n",
    "    snes.getKSP().setType(\"preonly\")\n",
    "    snes.getKSP().getPC().setType(\"lu\")\n",
    "    snes.getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    snes.setObjective(problem.obj)\n",
    "    snes.setFunction(problem.F, F_vec)\n",
    "    snes.setJacobian(problem.J, J=J_mat, P=None)\n",
    "    snes.setMonitor(lambda _, it, residual: print(it, residual))\n",
    "    up_copy = up.vector.copy()\n",
    "    snes.solve(None, up_copy)\n",
    "    problem.update_solution(up_copy)  # TODO can this be safely removed?\n",
    "    return up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "up_m = run_monolithic()\n",
    "(u_m, p_m) = (up_m.sub(0).collapse(), up_m.sub(1).collapse())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_vector_field_plot(mesh, vector_field, name, factor):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    values = np.zeros((mesh.geometry.x.shape[0], 3))\n",
    "    values[:, :2] = vector_field.compute_point_values()\n",
    "    grid.point_arrays[name] = values\n",
    "    grid.set_active_vectors(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    glyphs = grid.glyph(orient=name, factor=factor)\n",
    "    plotter.add_mesh(glyphs)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, u_m, \"u\", factor=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p_m, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Block FEniCSx formulation using a two independent function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_block():\n",
    "    # Function spaces\n",
    "    V = FunctionSpace(mesh, V_element)\n",
    "    Q = FunctionSpace(mesh, Q_element)\n",
    "\n",
    "    # Test and trial functions\n",
    "    (v, q) = (TestFunction(V), TestFunction(Q))\n",
    "    (du, dp) = (TrialFunction(V), TrialFunction(Q))\n",
    "    (u, p) = (Function(V), Function(Q))\n",
    "\n",
    "    # Variational forms\n",
    "    F = [nu * inner(grad(u), grad(v)) * dx + inner(grad(u) * u, v) * dx - div(v) * p * dx,\n",
    "         div(u) * q * dx]\n",
    "    J = [[derivative(F[0], u, du), derivative(F[0], p, dp)],\n",
    "         [derivative(F[1], u, du), derivative(F[1], p, dp)]]\n",
    "\n",
    "    # Boundary conditions\n",
    "    u_in = Function(V)\n",
    "    u_in.interpolate(u_in_eval)\n",
    "    u_wall = Function(V)\n",
    "    u_wall.interpolate(u_wall_eval)\n",
    "    bdofs_V_1 = locate_dofs_topological(V, mesh.topology.dim - 1, boundaries_1)\n",
    "    bdofs_V_2 = locate_dofs_topological(V, mesh.topology.dim - 1, boundaries_2)\n",
    "    inlet_bc = DirichletBC(u_in, bdofs_V_1)\n",
    "    wall_bc = DirichletBC(u_wall, bdofs_V_2)\n",
    "    bc = [inlet_bc, wall_bc]\n",
    "\n",
    "    # Class for interfacing with SNES\n",
    "    class NavierStokesProblem(object):\n",
    "        def __init__(self, F, J, solutions, bcs, P=None):\n",
    "            self._F = F\n",
    "            self._J = J\n",
    "            self._obj_vec = create_vector_block(F)\n",
    "            self._solutions = solutions\n",
    "            self._bcs = bcs\n",
    "            self._P = P\n",
    "\n",
    "        def update_solutions(self, x):\n",
    "            x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)\n",
    "            with BlockVecSubVectorWrapper(x, [V.dofmap, Q.dofmap]) as x_wrapper:\n",
    "                for x_wrapper_local, sub_solution in zip(x_wrapper, self._solutions):\n",
    "                    with sub_solution.vector.localForm() as sub_solution_local:\n",
    "                        sub_solution_local[:] = x_wrapper_local\n",
    "\n",
    "        def obj(self, snes, x):\n",
    "            self.F(snes, x, self._obj_vec)\n",
    "            return self._obj_vec.norm()\n",
    "\n",
    "        def F(self, snes, x, F_vec):\n",
    "            self.update_solutions(x)\n",
    "            with F_vec.localForm() as F_vec_local:\n",
    "                F_vec_local.set(0.0)\n",
    "            assemble_vector_block(F_vec, self._F, self._J, self._bcs, x0=x, scale=-1.0)\n",
    "\n",
    "        def J(self, snes, x, J_mat, P_mat):\n",
    "            J_mat.zeroEntries()\n",
    "            assemble_matrix_block(J_mat, self._J, self._bcs, diagonal=1.0)\n",
    "            J_mat.assemble()\n",
    "            if self._P is not None:\n",
    "                P_mat.zeroEntries()\n",
    "                assemble_matrix_block(P_mat, self._P, self._bcs, diagonal=1.0)\n",
    "                P_mat.assemble()\n",
    "\n",
    "    # Create problem\n",
    "    problem = NavierStokesProblem(F, J, (u, p), bc)\n",
    "    F_vec = create_vector_block(F)\n",
    "    J_mat = create_matrix_block(J)\n",
    "\n",
    "    # Solve\n",
    "    snes = PETSc.SNES().create(mesh.mpi_comm())\n",
    "    snes.setTolerances(max_it=20)\n",
    "    snes.getKSP().setType(\"preonly\")\n",
    "    snes.getKSP().getPC().setType(\"lu\")\n",
    "    snes.getKSP().getPC().setFactorSolverType(\"mumps\")\n",
    "    snes.setObjective(problem.obj)\n",
    "    snes.setFunction(problem.F, F_vec)\n",
    "    snes.setJacobian(problem.J, J=J_mat, P=None)\n",
    "    snes.setMonitor(lambda _, it, residual: print(it, residual))\n",
    "    solution = create_vector_block(F)\n",
    "    snes.solve(None, solution)\n",
    "    problem.update_solutions(solution)  # TODO can this be safely removed?\n",
    "    return (u, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(u_b, p_b) = run_block()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, u_b, \"u\", factor=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p_b, \"p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error computation between mixed and block cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_error(u_m, p_m, u_b, p_b):\n",
    "    u_m_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_m), grad(u_m)) * dx), op=MPI.SUM))\n",
    "    err_u_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_b - u_m), grad(u_b - u_m)) * dx), op=MPI.SUM))\n",
    "    p_m_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(p_m, p_m) * dx), op=MPI.SUM))\n",
    "    err_p_norm = np.sqrt(\n",
    "        mesh.mpi_comm().allreduce(assemble_scalar(inner(p_b - p_m, p_b - p_m) * dx), op=MPI.SUM))\n",
    "    print(\"Relative error for velocity component is equal to\", err_u_norm / u_m_norm)\n",
    "    print(\"Relative error for pressure component is equal to\", err_p_norm / p_m_norm)\n",
    "    assert np.isclose(err_u_norm / u_m_norm, 0., atol=1.e-10)\n",
    "    assert np.isclose(err_p_norm / p_m_norm, 0., atol=1.e-10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_error(u_m, p_m, u_b, p_b)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
